import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from tools.cfg import py2cfg
import os
import torch
from torch import nn
import cv2
import numpy as np
import argparse
from pathlib import Path
from tools.metric import Evaluator
from pytorch_lightning.loggers import CSVLogger
import random
import tqdm

def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True


def get_args():
    parser = argparse.ArgumentParser()
    arg = parser.add_argument
    arg("-c", "--config_path", type=Path, help="Path to the config.", required=True)
    return parser.parse_args()


class Supervision_Train:
    def __init__(self, config):
        self.config = config
        self.net = config.net
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

        self.loss = config.loss

        self.metrics_train = Evaluator(num_class=config.num_classes)
        self.metrics_val = Evaluator(num_class=config.num_classes)

        self.optimizer = config.optimizer
        self.lr_scheduler = config.lr_scheduler

        self.train_loader = config.train_loader
        self.val_loader = config.val_loader
        self.net.to(self.device)

    def forward(self, x):
        seg_pre = self.net(x)
        return seg_pre

    def train_one_epoch(self, epoch):
        self.net.train()
        for batch_idx, batch in tqdm(enumerate(self.train_loader)):
            img, mask = batch['img'].to(self.device), batch['gt_semantic_seg'].to(self.device)
            prediction = self.net(img)
            loss = self.loss(prediction, mask)
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            if self.config.use_aux_loss:
                pre_mask = nn.Softmax(dim=1)(prediction[0])
            else:
                pre_mask = nn.Softmax(dim=1)(prediction)
            pre_mask = pre_mask.argmax(dim=1)
            for i in range(mask.shape[0]):
                self.metrics_train.add_batch(mask[i].cpu().numpy(), pre_mask[i].cpu().numpy())


        self.on_train_epoch_end()

    def on_train_epoch_end(self):
        if 'vaihingen' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        elif 'potsdam' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        elif 'whubuilding' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        elif 'massbuilding' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        elif 'cropland' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        else:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union())
            F1 = np.nanmean(self.metrics_train.F1())

        OA = np.nanmean(self.metrics_train.OA())
        iou_per_class = self.metrics_train.Intersection_over_Union()
        eval_value = {'mIoU': mIoU,
                      'F1': F1,
                      'OA': OA}
        print('train:', eval_value)

        iou_value = {}
        for class_name, iou in zip(self.config.classes, iou_per_class):
            iou_value[class_name] = iou
        print(iou_value)
        self.metrics_train.reset()

    def validate_one_epoch(self):
        self.net.eval()
        with torch.no_grad():
            for batch_idx, batch in enumerate(self.val_loader):
                img, mask = batch['img'].to(self.device), batch['gt_semantic_seg'].to(self.device)
                prediction = self.forward(img)
                pre_mask = nn.Softmax(dim=1)(prediction)
                pre_mask = pre_mask.argmax(dim=1)
                for i in range(mask.shape[0]):
                    self.metrics_val.add_batch(mask[i].cpu().numpy(), pre_mask[i].cpu().numpy())

                loss_val = self.loss(prediction, mask)

        self.on_validation_epoch_end()

    def on_validation_epoch_end(self):
        if 'vaihingen' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        elif 'potsdam' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        elif 'whubuilding' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        elif 'massbuilding' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        elif 'cropland' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        else:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union())
            F1 = np.nanmean(self.metrics_val.F1())

        OA = np.nanmean(self.metrics_val.OA())
        iou_per_class = self.metrics_val.Intersection_over_Union()

        eval_value = {'mIoU': mIoU,
                      'F1': F1,
                      'OA': OA}
        print('val:', eval_value)
        iou_value = {}
        for class_name, iou in zip(self.config.classes, iou_per_class):
            iou_value[class_name] = iou
        print(iou_value)

        self.metrics_val.reset()

    def train(self, num_epochs):
        for epoch in range(num_epochs):
            self.train_one_epoch(epoch)
            self.validate_one_epoch()
            self.lr_scheduler.step()


# training
def main():
    args = get_args()
    config = py2cfg(args.config_path)
    seed_everything(42)

    # checkpoint_callback = ModelCheckpoint(save_top_k=config.save_top_k, monitor=config.monitor,
    #                                       save_last=config.save_last, mode=config.monitor_mode,
    #                                       dirpath=config.weights_path,
    #                                       filename=config.weights_name)
    # logger = CSVLogger('lightning_logs', name=config.log_name)

    model = Supervision_Train(config)
    model.train(config.max_epoch)
    # if config.pretrained_ckpt_path:
    #     model = Supervision_Train.load_from_checkpoint(config.pretrained_ckpt_path, config=config)

    # trainer = pl.Trainer(devices=config.gpus, max_epochs=config.max_epoch, accelerator='auto',
    #                      check_val_every_n_epoch=config.check_val_every_n_epoch,
    #                      callbacks=[checkpoint_callback], strategy='auto',
    #                      logger=logger)
    # trainer.fit(model=model, ckpt_path=config.resume_ckpt_path)


if __name__ == "__main__":
   main()
